import unittest
from pathlib import Path
import pandas as pd
from io import BytesIO
import tempfile
import shutil

from main import FontManager, RiskCalculator, ReportGenerator


class TestRiskCalculator(unittest.TestCase):
    """测试风险计算类"""
    
    @classmethod
    def setUpClass(cls):
        """设置测试数据"""
        cls.calculator = RiskCalculator()
        cls.sample_data = {
            '样本编号': 'TEST001',
            '姓名': '测试样本',
            '性别': '男',
            'rs1048943': 'CC',
            'rs25487': 'TT',
            'rs6682925': 'CT',
            'rs1042522': 'CC',
            '主检人': '张三',
            '审核人': '李四'
        }

    def test_get_snp_risk_level(self):
        """测试SNP风险等级计算"""
        risk_levels = self.calculator.get_snp_risk_level(self.sample_data)
        self.assertIsInstance(risk_levels, dict)
        for snp in ['rs1048943', 'rs25487', 'rs6682925', 'rs1042522']:
            self.assertIn(risk_levels[snp], ['风险', '正常'])

    def test_calculate_item_risk(self):
        """测试项目风险计算"""
        risks = self.calculator.calculate_item_risk(self.sample_data)
        self.assertIsInstance(risks, dict)
        for risk_score, risk_level, odds in risks.values():
            self.assertIsInstance(risk_score, float)
            self.assertIsInstance(risk_level, int)
            self.assertIsInstance(odds, float)
            self.assertGreaterEqual(risk_level, 1)
            self.assertLessEqual(risk_level, 4)
            self.assertLessEqual(odds, 3)


class TestReportGenerator(unittest.TestCase):
    """测试报告生成类"""
    
    @classmethod
    def setUpClass(cls):
        """设置测试环境"""
        cls.generator = ReportGenerator()
        cls.sample_data = {
            '样本编号': 'TEST001',
            '姓名': '测试样本',
            '性别': '男',
            'rs1048943': 'CC',
            'rs25487': 'TT',
            'rs6682925': 'CT',
            'rs1042522': 'GC',
            '主检人': '张三',
            '审核人': '李四'
        }
        cls.temp_dir = tempfile.mkdtemp()

    @classmethod
    def tearDownClass(cls):
        """清理测试环境"""
        shutil.rmtree(cls.temp_dir)

    def test_create_bar_chart(self):
        """测试柱状图创建"""
        risks = self.generator.risk_calculator.calculate_item_risk(self.sample_data)
        items = ['结直肠癌', '肺癌', '肝癌']  # 示例项目
        image_io = self.generator._create_bar_chart(items, risks)
        self.assertIsInstance(image_io, BytesIO)
        self.assertGreater(len(image_io.getvalue()), 0)

    def test_generate_report(self):
        """测试报告生成"""
        try:
            self.generator.generate_report(self.sample_data, self.temp_dir)
            expected_file = Path(self.temp_dir) / f"{self.sample_data['样本编号']}-{self.sample_data['姓名']}.docx"
            self.assertTrue(expected_file.exists())
            self.assertGreater(expected_file.stat().st_size, 0)
        except Exception as e:
            self.fail(f"生成报告失败: {str(e)}")

    def test_invalid_gender(self):
        """测试无效性别处理"""
        invalid_sample = self.sample_data.copy()
        invalid_sample['性别'] = '未知'
        with self.assertRaises(ValueError):
            self.generator.generate_report(invalid_sample, self.temp_dir)


def run_tests():
    """运行所有测试"""
    unittest.main(verbosity=2)


if __name__ == '__main__':
    run_tests()
